# -*- coding: utf-8 -*-
import numpy as np
import copy
from .node import Node
from .ElementBase import elementBase
from .tools import *
from .Loads import *
from .MATRIX_INV_DET import matrix_inv
class ps4(elementBase):
    """
    平面应力单元, 4节点,线性,节点自由度:Ux,U
    
    单元节点位移: [U1x,U1y,U2x,U2y,U3x,U3y,U4x,U4y]^T
    
    平面四边形单元,4节点线性分布,节点自由度:ux,uy;节点排序和abaqus的ps4一致;
    应力分量排序为 {11,22,33,12,23,13} ,和abaqus/explicit一致
    
    单元定义需: 节点坐标,模量E,泊松比nu,密度rho,厚度thickness
    
    支持的单元载荷: 体积力(N/m^3);压力载荷(N/m^2);表面力载荷(N/m^2)
    
    支持集中节点力加载:Fx,Fy
    """
    VTK_TYPE=9
    # 单元形函数
    N1=lambda xiVal,etaVal:0.25*(1-xiVal)*(1-etaVal)
    N2=lambda xiVal,etaVal:0.25*(1+xiVal)*(1-etaVal)
    N3=lambda xiVal,etaVal:0.25*(1+xiVal)*(1+etaVal)
    N4=lambda xiVal,etaVal:0.25*(1-xiVal)*(1+etaVal)
    
    # 形函数对xi的偏导
    N1xi=lambda xiVal,etaVal:0.25*(etaVal-1.0)
    N2xi=lambda xiVal,etaVal:0.25*(1.0-etaVal)
    N3xi=lambda xiVal,etaVal:0.25*(etaVal+1.0)
    N4xi=lambda xiVal,etaVal:0.25*(-1.0-etaVal)
    
    # 形函数对eta的偏导
    N1eta=lambda xiVal,etaVal:0.25*(xiVal-1)
    N2eta=lambda xiVal,etaVal:0.25*(-1-xiVal)
    N3eta=lambda xiVal,etaVal:0.25*(1+xiVal)
    N4eta=lambda xiVal,etaVal:0.25*(1-xiVal)
    
    def __init__(self,label:int,
                    node1:Node,
                    node2:Node,
                    node3:Node,
                    node4:Node,
                    rho=0.0,E=.0,
                    nu=.0,thickness=.0):
        super().__init__(label,2)
        self.node1 = node1
        self.node2 = node2
        self.node3 = node3
        self.node4 = node4
        self.rho = rho
        self.E = E
        self.nu = nu
        self.thickness = thickness
        # 2*2 高斯积分点自然坐标(xi_i,eta_i)和权重
        t=np.sqrt(1.0/3.0)
        self.GaussPoints=((t,t),(-t,t),(-t,-t),(t,-t))
        self.GaussWeights=(1,1,1,1)
    @property
    def get_connectivity(self):
        """
        返回单元的节点编号
        
        return: 
            List[int] : 单元的节点编号
        """
        return [self.node1.label,self.node2.label,self.node3.label,self.node4.label]
    @property
    def D(self):
        """
        平面应力状态,计算弹性变形矩阵
        
        return:
            np.ndarray : 弹性变形矩阵(3*3)
        """
        E,nu=self.E,self.nu
        return np.array([[1,nu,0],
                    [nu,1,0],
                    [0,0,(1 -nu)/2]])*E/(1-nu**2)
    def ShapeFun_at_xi_eta(self,xi,eta):
        """
        计算形函数N(2x8 matrix)在(xi,eta)处的值
        
        return:
            np.ndarray : 形函数N(2x8 matrix)在(xi,eta)处的值
        """
        xiVal,etaVal=xi,eta
        N1=ps4.N1(xiVal,etaVal)
        N2=ps4.N2(xiVal,etaVal)
        N3=ps4.N3(xiVal,etaVal)
        N4=ps4.N4(xiVal,etaVal)
        return np.array([[N1,0,N2,0,N3,0,N4,0],
                        [0,N1,0,N2,0,N3,0,N4]])
    def J_at_xi_eta(self,xi,eta):
        """
        计算在(xi,eta)处的雅克比矩阵
        
        return:
            np.ndarray(2,2) : 在(xi,eta)处的雅克比矩阵J
        """
        xiVal,etaVal=xi,eta
        x_xi =  ps4.N1xi(xiVal,etaVal)*self.node1.x+ps4.N2xi(xiVal,etaVal)*self.node2.x+ps4.N3xi(xiVal,etaVal)*self.node3.x+ps4.N4xi(xiVal,etaVal)*self.node4.x
        x_eta = ps4.N1eta(xiVal,etaVal)*self.node1.x+ps4.N2eta(xiVal,etaVal)*self.node2.x+ps4.N3eta(xiVal,etaVal)*self.node3.x+ps4.N4eta(xiVal,etaVal)*self.node4.x
        y_xi =  ps4.N1xi(xiVal,etaVal)*self.node1.y+ps4.N2xi(xiVal,etaVal)*self.node2.y+ps4.N3xi(xiVal,etaVal)*self.node3.y+ps4.N4xi(xiVal,etaVal)*self.node4.y
        y_eta = ps4.N1eta(xiVal,etaVal)*self.node1.y+ps4.N2eta(xiVal,etaVal)*self.node2.y+ps4.N3eta(xiVal,etaVal)*self.node3.y+ps4.N4eta(xiVal,etaVal)*self.node4.y
        return np.array([[x_xi,y_xi],[x_eta,y_eta]])
    def Jinv_at_xi_eta(self,xi,eta):
        """
        雅克比矩阵的逆矩阵,在(xi,eta)
        
        return:
            np.ndarray(2,2) : 雅克比矩阵的逆矩阵
        """
        J=self.J_at_xi_eta(xi,eta)
        return np.linalg.inv(J)
        # return matrix_inv(J)
    def B_at_xi_eta(self,xi,eta):
        """
        计算在(xi,eta)处的B矩阵(3*8)
        
        return:
            np.ndarray(3,8) : 在(xi,eta)处的B矩阵
        """
        Jinv=self.Jinv_at_xi_eta(xi,eta)
        if Jinv.shape!=(2,2):
            raise ValueError('Jinv shape should be (2,2)!')
        N1x=Jinv[0,0]*ps4.N1xi(xi,eta)+Jinv[0,1]*ps4.N1eta(xi,eta)
        N1y=Jinv[1,0]*ps4.N1xi(xi,eta)+Jinv[1,1]*ps4.N1eta(xi,eta)
        b1=np.array([[N1x,0],[0,N1y],[N1y,N1x]])
        
        N2x=Jinv[0,0]*ps4.N2xi(xi,eta)+Jinv[0,1]*ps4.N2eta(xi,eta)
        N2y=Jinv[1,0]*ps4.N2xi(xi,eta)+Jinv[1,1]*ps4.N2eta(xi,eta)
        b2=np.array([[N2x,0],[0,N2y],[N2y,N2x]])
        
        N3x=Jinv[0,0]*ps4.N3xi(xi,eta)+Jinv[0,1]*ps4.N3eta(xi,eta)
        N3y=Jinv[1,0]*ps4.N3xi(xi,eta)+Jinv[1,1]*ps4.N3eta(xi,eta)
        b3=np.array([[N3x,0],[0,N3y],[N3y,N3x]])
        
        N4x=Jinv[0,0]*ps4.N4xi(xi,eta)+Jinv[0,1]*ps4.N4eta(xi,eta)
        N4y=Jinv[1,0]*ps4.N4xi(xi,eta)+Jinv[1,1]*ps4.N4eta(xi,eta)
        b4=np.array([[N4x,0],[0,N4y],[N4y,N4x]])
        
        B=np.hstack((b1,b2,b3,b4))
        if B.shape!=(3,8):
            print(f"B({B.shape})=\n{B}")
            raise ValueError('B shape should be (3,8)!')
        return B
    @property
    def Ke(self):
        return self._Ke
    @property
    def Me(self,**kwargs):
        return self._Me
    @property
    def _Ke(self):
        """
        利用2*2高斯积分计算单元刚度矩阵Ke
        
        return:
            np.ndarray(8,8) : 单元刚度矩阵Ke
        """
        ke=np.zeros((8,8))
        for p,weight in zip(self.GaussPoints,self.GaussWeights):
            xi,eta=p
            B=self.B_at_xi_eta(xi,eta)
            jac=self.J_at_xi_eta(xi,eta)
            detJ=np.linalg.det(jac)
            if np.abs(detJ)<1e-10:
                print('Jacobian is too small, det(J)={:.2e}'.format(detJ))
                print(f"detJ={detJ}")
                continue
            Bt=B.T
            D=self.D
            ke+=self.thickness*weight*detJ*np.dot(np.dot(Bt,D),B)
        return ke.astype(float)
    @property
    def _Me(self,**kwargs):
        """
        计算单元质量矩阵Me,全局坐标系下,默认是集中质量矩阵,如果consistentMassMatrix=True,则是一致质量矩阵
        
        return:
            np.ndarray(8,8) : 单元质量矩阵Me
        """
        # lumped mass matrix
        def Triarean(node1,node2,node3):
            # 定义顶点坐标
            A = node1.coordinate
            B = node2.coordinate
            C = node3.coordinate
            # 计算向量 AB 和 AC
            AB = B - A
            AC = C - A
            # 计算叉乘
            cross_product = np.cross(AB, AC)
            # 计算叉乘的模
            cross_product_norm = np.linalg.norm(cross_product)
            # 计算面积
            return 0.5 * cross_product_norm
        # 一致单元质量矩阵
        if kwargs['consistentMassMatrix']:
            me=np.zeros((8,8))
            for p,weight in zip(self.GaussPoints,self.GaussWeights):
                xi,eta=p
                N=self.ShapeFun_at_xi_eta(xi,eta)
                jac=self.J_at_xi_eta(xi,eta)
                detJ=np.linalg.det(jac)
                if np.abs(detJ)<1e-10:
                    print('Jacobian is too small, det(J)={:.2e}'.format(detJ))
                    continue
                me+=self.thickness*weight*detJ*np.dot(N.T,N)
        else:
            me=0.25*self.rho*self.thickness*(Triarean(self.node1,self.node2,self.node3)+Triarean(self.node3,self.node4,self.node1))*np.eye(8) 
        return me.astype(float)
    @property
    def dofNum(self):
        """返回单元自由度数
        
        return:
            int : 单元自由度数,8
        """
        return 8
    @property
    def NodesList(self):
        """
        返回单元节点对象列表
        
        return:
            List[Node] : 单元节点对象列表
        """
        return[self.node1,self.node2,self.node3,self.node4]
    @property
    def ElemDofIndexs(self):
        """
        单元自由度索引列表
        
        return:
            List[int] : 单元自由度索引列表,0-based
        """
        return dofInds(self.node1.label,2)+dofInds(self.node2.label,2)+dofInds(self.node3.label,2)+dofInds(self.node4.label,2)
    @property
    def Face1Normal(self):
        """
        计算单元面1(1-2)外法向量
        
        return:
            np.ndarray(3,) : 单元面1(1-2)外法向量"""
        vec12=np.array([self.node2.x-self.node1.x,self.node2.y-self.node1.y])
        # 顺时针转
        return rotate_90(vec12,clockwise=True)
    @property
    def Face2Normal(self):
        """单元面2(2-3)外法向量,return(3,) array"""
        vec23=np.array([self.node3.x-self.node2.x,self.node3.y-self.node2.y])
        # 顺时针转
        return rotate_90(vec23,clockwise=True)
    @property
    def Face3Normal(self):
        """单元面3(3-4)外法向量,return(3,) array"""
        vec34=np.array([self.node4.x-self.node3.x,self.node4.y-self.node3.y])
        # 顺时针转
        return rotate_90(vec34,clockwise=True)
    @property
    def Face4Normal(self):
        """单元面4(4-1)外法向量,return(3,) array"""
        vec41=np.array([self.node1.x-self.node4.x,self.node1.y-self.node4.y])
        # 顺时针转
        return rotate_90(vec41,clockwise=True)
    def ProcessELoad2EqNodeForce(self,load):
        """
        计算单元载荷对应的等效节点力
        
        return:
            np.ndarray(8,1) : 单元载荷对应的等效节点力
        """
        if isinstance(load,BodyLoad):
            return self.__ProcessELoad2EqNodeForceBodyLoad(load)
        if isinstance(load,PressureLoad):
            return self.__ProcessELoad2EqNodeForcePressureLoad(load)
        if isinstance(load,SurfaceLoad):
            return self.__ProcessELoad2EqNodeForceSurfaceLoad(load)
        raise ValueError(f"ps4 do not support {load.__class__.__name__} load !")
    def __ProcessELoad2EqNodeForceBodyLoad(self,load):
        """
        将单元体力加载转换为等效节点力
        
        return:
            np.ndarray(8,1) : 单元体力对应的等效节点力
        """
        # 体力(N/m^3)
        eqforce=np.zeros((8,1),dtype=float)
        p=copy.deepcopy(load)
        bodyForceArray=np.zeros((2,1),dtype=float)
        if p.Direction=="x":
            bodyForceArray[0,0]=p.LoadValue
        if p.Direction=="y":
            bodyForceArray[1,0]=p.LoadValue
        # 高斯数值积分
        for point,weight in zip(self.GaussPoints,self.GaussWeights):
            xi,eta=point
            N=self.ShapeFun_at_xi_eta(xi,eta)
            jac=self.J_at_xi_eta(xi,eta)
            detJ=np.linalg.det(jac)
            eqforce+=self.thickness*weight*detJ*np.dot(N.T,bodyForceArray)
        if eqforce.shape!=(8,1):
            raise ValueError('eqforce shape should be (8,1)!')
        return eqforce
    def __ProcessELoad2EqNodeForcePressureLoad(self,load):
        """
        将压力分布载荷(N/m^2)转换为等效节点力
        
        return:
            np.ndarray(8,1) : 单元压力分布对应的等效节点力
        """
        # 压力分布载荷,计算公式见<<有限元方法基础教程>>p327
        eqforce=np.zeros((8,1))
        p=copy.deepcopy(load)
        if self.label not in p.get_ElemsList:
            return eqforce
        # 找出这个单元标签的位置在load.get_ElemsList中的所有位置
        positions = np.where(np.in1d(p.get_ElemsList, [self.label]))[0]
        # load定义中,可以一个单元的多个面都有压力,所以需要遍历
        for pos in positions:
            # 计算表面压力
            if p.FaceIdList[pos]==1:
                # 面1为作用面
                eqforce+=self.FaceLoadIntegral(self.Face1Normal,p.PressureValue,1)
            if p.FaceIdList[pos]==2:
                # 面2为作用面
                eqforce+=self.FaceLoadIntegral(self.Face2Normal,p.PressureValue,2)
            if p.FaceIdList[pos]==3:
                # 面3为作用面
                eqforce+=self.FaceLoadIntegral(self.Face3Normal,p.PressureValue,3)
            if p.FaceIdList[pos]==4:
                # 面1为作用面
                eqforce+=self.FaceLoadIntegral(self.Face4Normal,p.PressureValue,4)
        return eqforce
    def __ProcessELoad2EqNodeForceSurfaceLoad(self,load):
        """将面载荷转换为等效节点力,自定义方向的面载荷(N/m^2)"""
        eqforce=np.zeros((8,1))
        p=copy.deepcopy(load)
        # 找出单元标签的位置
        positions = np.where(np.in1d(p.ElementLabels, [self.label]))[0]
        for pos in positions:
            # 计算表面载荷等效节点力
            eqforce+=self.FaceLoadIntegral(p.DirectionVector,p.LoadValue,p.FaceIdList[pos])
        return eqforce
    def FaceLoadIntegral(self,DLoadDirection:np.ndarray,DloadValue,FaceId:int):
        """计算面载荷的线积分"""
        if not len(DLoadDirection)==3:
            raise ValueError('DLoadDirection should be a 3-dim vector!')
        px,py,_=DLoadDirection*DloadValue
        Ts=np.array([[px],
                    [py]])
        # 使用两点高斯积分进行line积分
        ps, ws = np.polynomial.legendre.leggauss(2)
        w1,w2=ws
        if FaceId==1:
            xis=ps
            etas=np.array([-1,-1])
            edgeLen=np.sqrt((self.node2.x-self.node1.x)**2+(self.node2.y-self.node1.y)**2)
        if FaceId==2:
            xis=np.array([1,1])
            etas=ps
            edgeLen=np.sqrt((self.node3.x-self.node2.x)**2+(self.node3.y-self.node2.y)**2)
        if FaceId==3:
            xis=ps
            etas=np.array([1,1])
            edgeLen=np.sqrt((self.node4.x-self.node3.x)**2+(self.node4.y-self.node3.y)**2)
        if FaceId==4:
            xis=np.array([-1,-1])
            etas=ps
            edgeLen=np.sqrt((self.node1.x-self.node4.x)**2+(self.node1.y-self.node4.y)**2)
        xi1,xi2=xis # type: ignore
        eta1,eta2=etas  # type: ignore
        return self.thickness*0.5*edgeLen*(w1*self.ShapeFun_at_xi_eta(xi=xi1,eta=eta1).T*Ts+w2*self.ShapeFun_at_xi_eta(xi=xi2,eta=eta2).T*Ts) # type: ignore

class PS4I(ps4):
    """
    使用wilson非协调模式进行shear locking修正的PS4单元
    
    单元公式参见: https://qinguoming.xyz/
    
    单元刚度矩阵已修正
    """
    # M1,M2为非协调模式的内部自由度的形函数
    M1=lambda xiVal,etaVal:  1-xiVal*xiVal
    M2=lambda xiVal,etaVal:  1-etaVal*etaVal
    
    # 对xi求偏导
    M1xi=lambda xiVal,etaVal: -2*xiVal+0.0*etaVal
    M2xi=lambda xiVal,etaVal: 0.0
    
    # 对eta求偏导
    M1eta=lambda xiVal,etaVal: 0.0
    M2eta=lambda xiVal,etaVal: -2*etaVal
    def __init__(self,label:int,
                    node1:Node,
                    node2:Node,
                    node3:Node,
                    node4:Node,
                    rho=0.0,E=.0,
                    nu=.0,thickness=.0):
        """
        初始化PS4I单元,使用wilson非协调模式进行shear locking修正
        """
        super().__init__(label,node1,node2,node3,node4,rho,E,nu,thickness)
        # 这里可以添加wilson非协调模式的相关参数和初始化
        
        
    def BI_at_xi_eta(self,xi,eta):
        """
        计算内部自由度贡献的应变矩阵BI在(xi,eta)的值
        
        return:
            np.ndarray(3,4) : 内部自由度贡献的应变矩阵BI
        """ 
        M1x,M1y=0,0
        M2x,M2y=0,0
        
        M1x,M1y=np.linalg.solve(self.J_at_xi_eta(xi,eta),
                        np.array([PS4I.M1xi(xi,eta),PS4I.M1eta(xi,eta)]))
        M2x,M2y=np.linalg.solve(self.J_at_xi_eta(xi,eta),
                        np.array([PS4I.M2xi(xi,eta),PS4I.M2eta(xi,eta)]))
        # 返回BI矩阵
        return np.array([[M1x,0,M2x,0],
                            [0,M1y,0,M2y],
                            [M1y,M1x,M2y,M2x]],dtype=float)
    @property
    def Ke(self):
        """
        计算wilson单元的刚度矩阵[Kcc,Kci; Kic,Kii]*{ue;alpha_e}={R;0}
        
        K=Kcc-Kci*inv(Kii)*Kic
        
        return:
            np.ndarray(8,8) : 单元刚度矩阵K
        """
        return self.Kcc-self.Kci@np.linalg.inv(self.Kii)@self.Kic
    @property
    def Kii(self):
        """
        计算内部自由度的刚度矩阵Kii
        
        return:
            np.ndarray(4,4) : 内部自由度的刚度矩阵Kii
        """
        Kii=np.zeros((4,4))
        for p,weight in zip(self.GaussPoints,self.GaussWeights):
            xi,eta=p
            BI=self.BI_at_xi_eta(xi,eta)
            jac=self.J_at_xi_eta(xi,eta)
            detJ=np.linalg.det(jac)
            if np.abs(detJ)<1e-10:
                print('Jacobian is too small, det(J)={:.2e}'.format(detJ))
                continue
            Kii+=self.thickness*weight*detJ*(BI.T@self.D@BI)
        return Kii.astype(float)
    @property
    def Kcc(self):
        """
        计算PS4单元的刚度矩阵Ke
        
        return:
            np.ndarray(8,8) : PS4单元的刚度矩阵Ke
        """
        return self._Ke
    @property
    def Kic(self):
        """
        计算PS4I单元的刚度矩阵Kic
        
        return:
            np.ndarray(4,8) : PS4单元的刚度矩阵Kic
        """
        Kic=np.zeros((4,8))
        for p,weight in zip(self.GaussPoints,self.GaussWeights):
            xi,eta=p
            B=self.B_at_xi_eta(xi,eta)
            BI=self.BI_at_xi_eta(xi,eta)
            jac=self.J_at_xi_eta(xi,eta)
            detJ=np.linalg.det(jac)
            if np.abs(detJ)<1e-10:
                print('Jacobian is too small, det(J)={:.2e}'.format(detJ))
                continue
            Kic+=self.thickness*weight*detJ*(BI.T@self.D@B)
        return Kic.astype(float)
    @property
    def Kci(self):
        """
        计算PS4I单元的刚度矩阵Kci
        
        return:
            np.ndarray(8,4) : PS4I单元的刚度矩阵Kci
        """
        Kci=np.zeros((8,4))
        for p,weight in zip(self.GaussPoints,self.GaussWeights):
            xi,eta=p
            B=self.B_at_xi_eta(xi,eta)
            BI=self.BI_at_xi_eta(xi,eta)
            jac=self.J_at_xi_eta(xi,eta)
            detJ=np.linalg.det(jac)
            if np.abs(detJ)<1e-10:
                print('Jacobian is too small, det(J)={:.2e}'.format(detJ))
                continue
            Kci+=self.thickness*weight*detJ*(B.T@self.D@BI)
        return Kci.astype(float)